package org.apache.commons.graph.algorithms.flow.internal

import org.apache.commons.graph.DirectedGraph
import org.apache.commons.graph.Mapper
import org.apache.commons.graph.algorithms.visit.BaseGraphVisitHandler
import org.apache.commons.graph.algorithms.visit.VisitState
import org.apache.commons.graph.weight.OrderedMonoid

/**
 * Provides standard operations for max-flow algorithms,
 * like Ford-Fulkerson or Edmonds-Karp.
 *
 * @param <V> the vertex type
 * @param <W> the weight type
</W></V> */
internal class FlowNetworkHandler<V, E, W>(
    private val flowNetwork: DirectedGraph<V, E>,
    private val source: V,
    private val target: V,
    private val weightOperations: OrderedMonoid<W>,
    private val weightedEdges: Mapper<E, W>
) : BaseGraphVisitHandler<V, E, DirectedGraph<V, E>, W>() {

    private var maxFlow: W

    private val residualEdgeCapacities: MutableMap<E, W> = HashMap()

    // these are new for each new visit of the graph
    private var predecessors: org.apache.commons.graph.algorithms.shortestpath.PredecessorsList<V, E, W>?

    private var foundAugmentingPath = false

    init {

        maxFlow = weightOperations.identity()

        for (edge in flowNetwork.edges) {
            residualEdgeCapacities[edge] = weightedEdges.map(edge)
        }

        predecessors = null
    }

    /**
     * {@inheritDoc}
     */
    override fun discoverEdge(head: V, edge: E, tail: V): VisitState {
        val residualEdgeCapacity = residualEdgeCapacities[edge]
        // avoid expanding the edge when it has no residual capacity
        if (weightOperations.compare(residualEdgeCapacity, weightOperations.identity()) <= 0) {
            return VisitState.SKIP
        }
        predecessors!!.addPredecessor(tail, head)
        return VisitState.CONTINUE
    }

    /**
     * {@inheritDoc}
     */
    override fun discoverGraph(graph: DirectedGraph<V, E>) {
        // reset ausiliary structures for a new graph visit
        predecessors = org.apache.commons.graph.algorithms.shortestpath.PredecessorsList(
            graph,
            weightOperations,
            weightedEdges
        )
        foundAugmentingPath = false
    }

    /**
     * {@inheritDoc}
     */
    override fun discoverVertex(vertex: V): VisitState {
        return finishVertex(vertex)
    }

    /**
     * {@inheritDoc}
     */
    override fun finishVertex(vertex: V): VisitState {
        if (vertex == target) {
            // search ends when target vertex is reached
            foundAugmentingPath = true
            return VisitState.ABORT
        }
        return VisitState.CONTINUE
    }

    /**
     * Checks whether there is an augmenting path in the flow network,
     * given the current residual capacities.
     * @return true if there is an augmenting path, false otherwise
     */
    fun hasAugmentingPath(): Boolean {
        return foundAugmentingPath
    }

    override fun onCompleted(): W? {
        return maxFlow
    }

    /**
     * Updates the residual capacities in the flow network,
     * based on the most recent augmenting path.
     */
    fun updateResidualNetworkWithCurrentAugmentingPath() {
        // build actual augmenting path
        val augmentingPath = predecessors!!.buildPath(source, target)

        // find flow increment
        var flowIncrement: W? = null
        for (edge in augmentingPath.edges) {
            val edgeCapacity = residualEdgeCapacities[edge]
            if (flowIncrement == null || weightOperations.compare(edgeCapacity, flowIncrement) < 0
            ) {
                flowIncrement = edgeCapacity
            }
        }

        // update max flow and capacities accordingly
        maxFlow = weightOperations.append(maxFlow, flowIncrement!!)
        for (edge in augmentingPath.edges) {
            // decrease capacity for direct edge
            val directCapacity = residualEdgeCapacities[edge]!!
            residualEdgeCapacities[edge] =
                weightOperations.append(directCapacity, weightOperations.inverse(flowIncrement))

            // increase capacity for inverse edge
            val vertexPair = flowNetwork.getVertices(edge)
            val inverseEdge = flowNetwork.getEdge(vertexPair!!.tail, vertexPair.head)!!
            val inverseCapacity = residualEdgeCapacities[inverseEdge]!!
            residualEdgeCapacities[inverseEdge] = weightOperations.append(inverseCapacity, flowIncrement)
        }
    }
}